package immutablecollections;

import immutablecollections.exceptions.*;
import immutablecollections.functions.*;
import immutablecollections.misc.*;
import immutablecollections.textboxes.*;
import java.io.*;
import java.util.*;

/**
 * An immutable binary tree.
 * 
 * ## Introduction
 * 
 * This class is the heart of the Immutable collections library.
 * It is intended to be used only as a component of the other classes.
 * 
 * An `ImTree` is an AVL tree (a balanced binary tree) where each node stores some arbitrary data.
 * 
 * Note that, in this class, there is no concept of the data that is being stored being Comparable.
 * This functionality is added by {@link ImSortedSet}.
 * 
 * A node in an `ImTree` is either a leaf node `Nil` or a `Node` that has two children that
 * are themselves `ImTrees`.
 * 
 * Each `Node` can contain a value of an arbitrary type and two Integers
 * representing the height and size of the tree rooted at that node. A `Nil` has no data.
 * 
 * Note that this definition does not, of itself, specify that the tree is balanced. We enforce that
 * invariant in each method that adds/removes nodes.
 * 
 * Consider an example tree with six non nil nodes:
 * 
 * <img src="doc-files/tree-abcdef.png" alt="tree-abcdef"  style="width: 20%"/></p>
 *     
 * If we show the nil nodes then it looks like this:
 * 
 * <img src="doc-files/tree-abcdef-with-nulls.png" alt="tree-abcdef-with-nulls"  style="width: 20%"/></p>
 *     
 * Each `Node` also has a `size` value defined as the sum of the sizes of its children plus one.
 * nil Nodes are considered to have a size of zero.
 * The size represents how many non nil nodes there are in the tree rooted at n.
 * 
 * <img src="doc-files/tree-abcdef-with-sizes.png" alt="tree-abcdef-with-sizes"  style="width: 20%"/></p>
 *
 * Each node also has a height value representing the size of the longest path from that node to a leaf
 * node.
 * 
 * The height of a node is the maximum of the heights of its children. Nil nodes are considered to
 * have a height of zero.
 * 
 * Let's annotate our example with the heights:
 * 
 * <img src="doc-files/tree-abcdef-with-heights.png" alt="tree-abcdef-with-heights"  style="width: 20%"/></p>
 *
 * Because the tree is balanced, this means that the heights of the children of a node will differ
 * by at most one.
 * 
 * Each node is considered to have a *rank* that represents its position in the tree in a pre-order scan.
 * 
 * Ranks start at one (exactly as Nature intended!)
 * 
 * Let's annotate our example with the ranks:
 * 
 * <img src="doc-files/tree-abcdef-with-ranks.png" alt="tree-abcdef-with-ranks"  style="width: 20%"/></p>
 *
 * We don't store the ranks. To calculate the rank of a node or to find a node at a particular rank we can use
 * the size of child nodes to derive the answer.
 * 
 * ## Don't go changing to try to please me...
 * 
 * These are **immutable** Collections so the `insert` and `remove` methods don't actually change existing trees.
 * Instead they create a new tree with a node added or deleted as appropriate, reusing as many of
 * the old nodes as possible.
 * 
 * ## References
 * 
 * # [Implementing Sets Efficiently in a Functional Language] Stephen Adams
 * 
 * [adams]: http://groups.csail.mit.edu/mac/users/adams/BB/92-10.ps
 * 
 * @see ImCollection
 */

class ImTree<A> implements Serializable
{
    // TODO replace with getters? Deal with nil?
    final private A element;
    final private ImTree<A> left;
    final private ImTree<A> right;

    final private int height;
    final private int size;
    final private int hashCode;

    static final ImTree<?> nil = new ImTree<Object>();
    private static final int MAX_SUBTREE_HEIGHT_DIFFERENCE_PLUS_ONE = 2;

    @SuppressWarnings("unchecked")
    public static <A> ImTree<A> Nil()
    {
        return (ImTree<A>) nil;
    }

    protected ImTree()
    {
        this(null, null, null, 0, 0);
    }

    private Object readResolve() throws ObjectStreamException
    {
        return getElement() == null
                ? nil
                : this;
    }

    protected ImTree(final A element, final ImTree<A> left, final ImTree<A> right, final int height, final int size)
    {
        this.element = element;
        this.left = left;
        this.right = right;
        this.height = height;
        this.size = size;
        this.hashCode = element == null
                ? 0
                : element.hashCode() + left.hashCode() + right.hashCode();

    }

    public ImTree(final A a, final ImTree<A> left, final ImTree<A> right)
    {
        this(a, left, right, 1 + Math.max(left.height, right.height), left.size + right.size + 1);
    }

    public static <A> ImTree<A> newLeaf(final A a)
    {
        return new ImTree<A>(a, ImTree.<A> Nil(), ImTree.<A> Nil());
    }

    public static <A> ImTree<A> on(Collection<A> elements)
    {
        ImTree<A> result = ImTree.Nil();

        int index = 1;
        for (A a : elements)
        {
            result = result.insert(index++, a);
        }

        return result;
    }

    private static enum Balance
    {
        left, right, balanced, unbalanced, leftUnbalanced, rightUnbalanced
    }

    static <A> ImTree<A> newBalancedTree(final A newA, final ImTree<A> newLeft, final ImTree<A> newRight)
    {

        Balance b = getBalance(newLeft, newRight);

        switch (b)
        {
        case left:
            return newTreeL(newA, newLeft, newRight);

        case right:
            return newTreeR(newA, newLeft, newRight);

        case balanced:
            return new ImTree<A>(newA, newLeft, newRight);

        }

        throw new InvalidStateException("difference in tree height is too great");
    }

    private static <A> Balance getBalance(final ImTree<A> left, final ImTree<A> right)
    {
        int diff = left.height - right.height;

        if (diff == MAX_SUBTREE_HEIGHT_DIFFERENCE_PLUS_ONE)
            return Balance.left;

        if (diff == -MAX_SUBTREE_HEIGHT_DIFFERENCE_PLUS_ONE)
            return Balance.right;

        if (Math.abs(diff) < MAX_SUBTREE_HEIGHT_DIFFERENCE_PLUS_ONE)
            return Balance.balanced;

        if (diff > MAX_SUBTREE_HEIGHT_DIFFERENCE_PLUS_ONE)
            return Balance.leftUnbalanced;

        if (diff < -MAX_SUBTREE_HEIGHT_DIFFERENCE_PLUS_ONE)
            return Balance.rightUnbalanced;

        throw new InvalidStateException("Unknown balance");

    }

    static <A> ImTree<A> newTreeL(final A element, final ImTree<A> b, final ImTree<A> c)
    {
        final ImTree<A> d = b.getLeft();
        final ImTree<A> e = b.getRight();

        /**
         * We are a. If our left node (b) is too tall then either d or e is the
         * culprit. It can't be both
         *          a    
         *      b       c
         *    d   e 
         *        
         */

        if (e.height > d.height)
        {
            /**
             * Ok - so it is e that is too tall
             * So - e cannot be nil - so it must have children f and g
             * 
             *  What it would        After swizzling
             *  be if we didn't
             *  swizzle
             * 
             *         a        1        e'   
             *     b       c    2      b'  a' 
             *   d   e          3     d f g c
             *      f g         4
             */
            return new ImTree<A>(e.getElement(), new ImTree<A>(b.getElement(), d, e.getLeft()), new ImTree<A>(element,
                    e.getRight(), c));
        }
        else
        {
            /**
             * Ok - so either d and e are the same height. They can't both be nil, otherwise b could not be
             * or d is taller than e (in which case d is not nil)
             * 
             * d and e might be nil
             * 
             *  What it would        After swizzling
             *  be if we didn't
             *  swizzle
             *  
             *      a         1       b'                 
             *    b   c       2     d   a'               
             *   d e          3        e c 
             *                
             */
            return new ImTree<A>(b.getElement(), d, new ImTree<A>(element, e, c));
        }

    }

    static <A> ImTree<A> newTreeR(final A element, final ImTree<A> b, final ImTree<A> c)
    {

        final ImTree<A> d = c.getLeft();
        final ImTree<A> e = c.getRight();

        /**
         * We are a. Our right node c is too tall.
         * So either d or e is the
         * culprit. It can't be both
         *          a
         *        /   \   
         *       b     c
         *            / \
         *           d   e 
         *        
         */

        if (d.height > e.height)
        {
            /**
             * Ok - so it is d that is too tall. We need to do a double rotation.
             * 
             *    What it would        After swizzling
             *    be if we didn't
             *    swizzle
             *    
             *         a         1       d'
             *       /   \              /  \   
             *      b     c      2     a'   c'
             *           / \          / \  / \
             *          d   e    3    b f g   e
             *         / \  
             *        f   g      4
             */
            return new ImTree<A>(d.getElement(), new ImTree<A>(element, b, d.getLeft()), new ImTree<A>(c.getElement(),
                    d.getRight(), e));
        }
        else
        {
            /**
             * Ok - so it is e that is too tall or d and e have the same height
             * 
             *     What it would        After swizzling
             *     be if we didn't
             *     swizzle
             *     
             *                    
             *          a         1       c'                
             *        /   \              /  \      
             *       b     c      2     a'   e
             *            / \          / \ 
             *           d   e    3   b   d 
             *     
             *
             */
            return new ImTree<A>(c.getElement(), new ImTree<A>(element, b, d), e);
        }
    }

    static <A> ImTree<A> merge(final ImTree<A> left, final ImTree<A> right)
    {
        // If the left child is nil then just return the right child (which could, itself be nil)
        if (left == nil)
            return right;

        // and vice versa
        if (right == nil)
            return left;

        /**
         * Uh oh - both children are non nil.
         */

        if (left.height >= right.height)
        {
            return concat3(left.getElement(), left.getLeft(), merge(left.getRight(), right));
        }
        else
        {
            return concat3(right.getElement(), merge(left, right.getLeft()), right.getRight());
        }
    }

    // end - Constructors and factory methods

    static <A> boolean isNil(final ImTree<A> treeToBeChecked)
    {
        return treeToBeChecked == nil;
    }

    // Shelf

    /**
     * Insert a new node with the value `elementToAdd` at index `indexStartingAtOne`
     * 
     * If we add node g to this example at index 3 we will generate a new tree with three new nodes
     * a', b' and g. None of the old nodes will be affected.
     * 
     *              a                       a'
     *            /   \                   /   \
     *           /     \                 /     \                   
     *          b       c  ----+        b'      |                
     *         /      /   \    |       / \      | 
     *        d      e     f   |      |   g     | 
     *        |                |      |         |
     *        +-----------------------+         |
     *                         |                |
     *                         +----------------+
     *                                          
     *                                          
     */
    public ImTree<A> insert(final int indexStartingAtOne, final A elementToAdd)
    {
        if (this == ImTree.nil)
            return ImTree.newLeaf(elementToAdd);

        final int localIndex = indexStartingAtOne - (getLeft().size + 1);

        return localIndex <= 0
                ? ImTree.newBalancedTree(getElement(), getLeft().insert(indexStartingAtOne, elementToAdd), getRight())
                : ImTree.newBalancedTree(getElement(), getLeft(), getRight().insert(localIndex, elementToAdd));
    }

    public static <A> ImTree<A> replaceAtIndex(final ImTree<A> tree, final int indexStartingAtOne, final A newElement)
    {
        final int localIndex = indexStartingAtOne - (tree.getLeft().size + 1);

        return localIndex == 0
                ? tree.getElement() == newElement
                        ? tree
                        : ImTree.newBalancedTree(newElement, tree.getLeft(), tree.getRight())
                : localIndex < 0
                        ? ImTree.newBalancedTree(tree.getElement(),
                                replaceAtIndex(tree.getLeft(), indexStartingAtOne, newElement), tree.getRight())
                        : ImTree.newBalancedTree(tree.getElement(), tree.getLeft(),
                                replaceAtIndex(tree.getRight(), localIndex, newElement));
    }

    public ImTree<A> getNodeAtIndex(final int indexStartingAtOne)
    {
        ImIndexOutOfBoundsException.check(indexStartingAtOne, size, "indexStartingAtOne");

        final int localIndex = indexStartingAtOne - (getLeft().size + 1);

        return localIndex == 0
                ? this
                : localIndex < 0
                        ? getLeft().getNodeAtIndex(indexStartingAtOne)
                        : getRight().getNodeAtIndex(localIndex);
    }

    /**
     * Return a tree that is the same as `this` but without the element
     * stored at the node with index `indexStartingAtOne`
     */
    public ImTree<A> remove(int indexStartingAtOne)
    {
        ImIndexOutOfBoundsException.check(indexStartingAtOne, size, "indexStartingAtOne");
        final int localIndex = indexStartingAtOne - (getLeft().size + 1);

        return localIndex == 0
                ? removeRoot()
                : localIndex < 0
                        ? ImTree.newBalancedTree(getElement(), getLeft().remove(indexStartingAtOne), getRight())
                        : ImTree.newBalancedTree(getElement(), getLeft(), getRight().remove(localIndex));
    }

    /**
     * The tree with its root removed.
     * 
     * We are removing a in the following diagrams
     * 
     *        a            =>               c
     *       / \                        
     *      -   c                        
     *                                     
     *        a            =>               b
     *       / \                         
     *      b   -                          
     *   
     *
     *        a            =>            merge(b,c)        
     *       / \                
     *      b   c              
     *                                             
     */
    public ImTree<A> removeRoot()
    {
        return merge(getLeft(), getRight());
    }

    // End Shelf

    // Printing
    @Override
    public String toString()
    {
        return isNil(this)
                ? "-"
                : "" + getElement() + (getHeight() > 1
                        ? " (" + getLeft() + " " + getRight() + ")"
                        : "");
    }

    public String toBoxString()
    {
        return toBox(this).toString();
    }

    protected static <A> TextBox toBox(final ImTree<A> tree)
    {
        if (tree == nil)
            return LeafTextBox.with("-");

        final String myText = tree.getElement().toString();

        if ((tree.getLeft() == nil) && (tree.getRight() == nil))
            return LeafTextBox.with(myText);

        final TextBox leftChildBox = toBox(tree.getLeft());
        final TextBox rightChildBox = toBox(tree.getRight());

        final int leftWidth = leftChildBox.getWidth();
        final int rightWidth = rightChildBox.getWidth();

        /**
         *     +--------+
         *     |        |
         *     +--------+
         *     +----+  +--------------+
         *     |    |  |              | 
         *     +----+  +--------------+ 
         */
        final int width = Math.max(myText.length(), leftWidth + 1 + rightWidth);

        final LeafTextBox gap = LeafTextBox.centred("", width - (leftWidth + rightWidth));
        final LeftRightBox children = LeftRightBox.with(leftChildBox, gap, rightChildBox);

        final String dots = TextUtils.repeat(".", (leftWidth + 1) / 2 + gap.getWidth() + (rightWidth + 1) / 2);
        final String spaceAndDots = TextUtils.repeat(" ", leftWidth / 2) + dots;

        return TopDownBox.with(LeafTextBox.centred(myText, width), LeafTextBox.with(spaceAndDots), children);
    }

    // end - Printing

    public String elementToString()
    {
        return this == nil
                ? "-"
                : getElement().toString();
    }

    /**
     * The name `concat3` is taken directly from [Implementing Sets Efficiently in a Functional Language][adams]
     * 
     * Consider an example
     * let's denote newBalancedTree() by BT() and concat3 by C3() (roughly)
     * 
     *       C3(f)
     *         / \        
     *       b    g           
     *     /   \         
     *    a     d   
     *         / \   
     *        c   e  
     * 
     * expanding C3:
     *        
     *     BT(b)       
     *     /   \         
     *    a  C3(f)
     *         / \   
     *        d   g
     *       / \
     *      c   e         
     *           
     * Expanding C3:
     *        
     *     BT(b)       
     *     /   \         
     *    a     f
     *         / \   
     *        d   g
     *       / \
     *      c   e
     *               
     * Expanding BT:
     *  
     *         d
     *       /   \ 
     *      b     f
     *     / \   / \
     *    a   c e   g
     *              
     */

    public static <A> ImTree<A> concat3(A element, ImTree<A> left, ImTree<A> right)
    {
        Balance b = getBalance(left, right);

        switch (b)
        {
        case left:
        case leftUnbalanced:
            return newBalancedTree(left.getElement(), left.getLeft(), concat3(element, left.getRight(), right));

        case right:
        case rightUnbalanced:
            return newBalancedTree(right.getElement(), concat3(element, left, right.getLeft()), right.getRight());

        case balanced:
            return new ImTree<A>(element, left, right);
        }

        throw new InvalidStateException("");
    }

    /**
     * 
     */
    public boolean isBalanced()
    {
        return this == nil
                ? true
                : getBalance(getLeft(), getRight()) == Balance.balanced && getLeft().isBalanced()
                        && getRight().isBalanced();
    }

    public ImList<A> toList()
    {
        ImTreeZipper<A> z = ImTreeZipper.onRightmost(this);

        if (z.getFocus() == ImTree.nil)
        {
            return (ImList<A>) ImList.empty();
        }

        ImList<A> result = ImList.on(z.getElement());
        ImMaybe<ImTreeZipper<A>> m = z.previous();

        while (m.hasValue())
        {
            ImTreeZipper<A> current = m.getValue();
            result = ImList.cons(current.getElement(), (ImList<A>) result);

            m = current.previous();
        }
        return result;
    }

    /**
     * The tree that has the same shape as `this` and where each element is the result of evaluating
     * the single argument function `fn` on
     * the corresponding element of `this`.
     */
    public <O> ImTree<O> map(Function1<O> fn)
    {
        return isNil(this)
                ? ImTree.<O> Nil()
                : new ImTree<O>(fn.invoke(getElement()), getLeft().map(fn), getRight().map(fn));
    }

    public A getElement()
    {
        return element;
    }

    /**
     * The left sub-tree of `this`.
     */
    public ImTree<A> getLeft()
    {
        return left;
    }

    /**
     * The right sub-tree of `this`.
     */
    public ImTree<A> getRight()
    {
        return right;
    }

    public int getHeight()
    {
        return height;
    }

    public int size()
    {
        return size;
    }

    @Override
    public int hashCode()
    {
        return hashCode;
    }

    /**
     * Get the rank of the root node of this tree in this tree.
     */
    public int getRank()
    {
        return isNil(this)
                ? 0
                : left.size + 1;
    }

}
